# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/9/8 上午10:45
# @Author  : shutao
# @FileName: cpnet_base.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

import os

import paddle

from modules.utils import LRScheduler
from modules.exp.base_exp import BaseExp


class Exp(BaseExp):
    """The experiment of cpnet."""

    def __init__(self):
        super().__init__()

        # -------------------- model config --------------------
        self.num_classes = 4
        self.out_size = (768, 768)
        self.pretrained = True
        self.prior_channels = 512
        self.am_kernel_size = 11
        self.proir_size = 64
        self.channels = 512
        self.backbone = "resnet18"

        # -------------------- dataloader config --------------------
        self.data_num_workers = 0
        self.create_datalist = False  # 是否增广并创建便于训练的数据集

        # -------------------- transform config --------------------

        # --------------------  training config --------------------
        self.max_iters = 60000
        self.max_epoch = 30
        self.scheduler = "PolynomialDecay"
        self.base_lr = 0.01
        self.power = 0.9

        self.weight_decay = 5e-4
        self.momentum = 0.9
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]

        # --------------------  testing config --------------------

    def get_model(self):
        """
        构造网络模型
        :return:
        """
        from ..models.heads import CPNet

        self.model = CPNet(
            n_classes=self.num_classes,
            out_size=self.out_size,
            pretrained=self.pretrained,
            prior_channels=self.prior_channels,
            prior_size=self.proir_size,
            am_kernel_size=self.am_kernel_size,
            channels=self.channels,
            backbone=self.backbone
        )
        return self.model

    def get_data_loader(self, batch_size, no_aug=False):
        """
        构造数据读取器
        :param batch_size:
        :param no_aug:
        :return:
        """
        from ..data.datasets import HouseMapDataset
        from ..data import image_augment
        from scripts import create_list_of_dataset

        if self.create_datalist:
            image_augment()  # 数据增广
            create_list_of_dataset()  # 创建并划分训练数据列表文件
        td = HouseMapDataset(img_size=self.out_size)
        td_loader = paddle.io.DataLoader(td, batch_size=batch_size, shuffle=True, drop_last=True)
        return td_loader, len(td_loader)

    def get_loss_function(self, preds, loss: paddle.Tensor, **kwargs) -> paddle.Tensor:

        pass

    def get_optimizer(self, scheduler, parameters):
        opt = paddle.optimizer.Momentum(learning_rate=scheduler,
                                        parameters=parameters, weight_decay=self.weight_decay, momentum=self.momentum)
        return opt

    def get_lr_scheduler(self, lr, iters_per_epoch, **kwargs):
        lr_scheduler = LRScheduler(self.scheduler,
                                   self.base_lr,
                                   max_iters=self.max_iters,
                                   power=self.power
                                   )
        scheduler = lr_scheduler.lr_func
        return scheduler

    def get_training_subject(self, data, loss_func):
        pass

    def get_evaluator(self, batch_size):
        from ..data.datasets import HouseMapDataset

        vd = HouseMapDataset(is_train=False, img_size=self.out_size)
        vd_loader = paddle.io.DataLoader(vd, batch_size=1, shuffle=False, drop_last=False)
        return vd_loader

    def eval(self, model, evaluator, weights):
        pass
